import torch
a = torch.rand(2,2,4)
a[0][0][0] = 0.4018
a[0][1][0] =  0.2399
a[1][0][0] = 0.0760
a[1][1][0] =  0.1233

a[0][0][1] = 0.1839
a[0][1][1] =   0.4173

a[1][0][1] = 0.2400
a[1][1][1] =  0.0497


a[0][0][2] = 0.9027
a[0][1][2] =  0.4909
a[1][0][2] = 0.9448
a[1][1][2] =  0.4893

a[0][0][0] = 1
a[0][0][1] = 0
a[0][0][2] = 1
a[0][0][3] = 0

a[1][0][0] = 0
a[1][0][1] = 1
a[1][0][2] = 0
a[1][0][3] = 1


print(a.size())
print(a[:,:,0])
print(a[:,:,1])
print(a[:,:,2])
b = torch.fft.fftn(a,s=3,dim=2)  # ok
print(b[0][0])
print(b[1][0])



# def schatten_norm(matrix,p=0.5):
#     _, sigma, _ = torch.svd(matrix, some=True)
#     # print(matrix)
#     # print(sigma)
#     return torch.pow(torch.sum(torch.pow(sigma,p)),1/p)
# print('----------')
# print(b[:,:,0])
# print(schatten_norm(b[:,:,0]))
# print(b[:,:,1])
# print(schatten_norm(b[:,:,1]))
# print(b[:,:,2])
# print(schatten_norm(b[:,:,2]))
# # print(b[:,:,3])
# # print(b[:,:,4])
# # print(b[:,:,5])